import torch
import torch.nn as nn


class Conv3dC1(nn.Module):
    def __init__(self, c1, c2, padding=(1, 1, 1), pool_kernel=(2, 2, 2), pool_stride=(2, 2, 2)):
        super(Conv3dC1, self).__init__()
        self.conv = nn.Conv3d(c1, c2, kernel_size=(3, 3, 3), padding=padding)
        self.bn = nn.BatchNorm3d(c2)
        self.relu = nn.LeakyReLU(negative_slope=0.1)
        self.pool = nn.MaxPool3d(kernel_size=pool_kernel, stride=pool_stride)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = self.pool(x)
        return x


class Conv3dC2(nn.Module):
    def __init__(self, c1, c2):
        super(Conv3dC2, self).__init__()
        self.conv1 = nn.Conv3d(c1, c2, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.bn1 = nn.BatchNorm3d(c2)
        self.conv2 = nn.Conv3d(c2, c2, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.bn2 = nn.BatchNorm3d(c2)
        self.relu = nn.LeakyReLU(negative_slope=0.1)
        self.pool = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.pool(x)
        return x


class C3D(nn.Module):
    def __init__(self, num_classes=5):
        super(C3D, self).__init__()
        self.conv1 = Conv3dC1(3, 12, padding=(1, 1, 1), pool_kernel=(1, 2, 2), pool_stride=(1, 2, 2))
        self.conv2 = Conv3dC1(12, 24, padding=(1, 1, 1), pool_kernel=(2, 2, 2), pool_stride=(2, 2, 2))
        self.conv3 = Conv3dC2(24, 48)
        self.conv4 = Conv3dC2(48, 96)
        self.conv5 = Conv3dC1(96, 192)
        self.fc1 = nn.Linear(13440, 1024)
        self.fc2 = nn.Linear(1024, num_classes)
        self.initweight()

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = x.view(-1, 13440)
        x = self.fc1(x)
        x = self.fc2(x)
        return x

    def initweight(self):
        for m in self.modules():
            if isinstance(m, (nn.Conv3d, nn.Linear)):
                nn.init.kaiming_normal_(m.weight)


# def init_param(model):
#     for m in model.modules():
#         if isinstance(m, (nn.Conv2d, nn.Linear)):
#             nn.init.kaiming_normal_(m.weight)


if __name__ == "__main__":
    inputs = torch.rand(1, 3, 31, 240, 320)
    net = C3D(num_classes=5)
    outputs = net.forward(inputs)
    print(outputs.size())
